import asyncio
import aiohttp
import json
import pandas as pd
import time
import argparse
import uuid
from typing import List, Dict
import random

# 配置参数
parser = argparse.ArgumentParser(description='测试网络信访件服务器并发处理能力')
parser.add_argument('--num_requests', type=int, default=10, help='并发请求数量')
parser.add_argument('--excel_file', type=str, default='wangluojian.xlsx', help='Excel文件路径')
parser.add_argument('--server_url', type=str, default='http://localhost:9000/api/getNetLetters', help='服务器URL')
args = parser.parse_args()

def prepare_test_data(row: pd.Series) -> Dict:
    """准备单个测试数据"""
    # 处理反映人信息
    fyrs = []
    if pd.notna(row["反映人姓名"]):
        fyr_names = str(row["反映人姓名"]).split(";")
        for name in fyr_names:
            if name.strip():
                fyrs.append({
                    "bh": str(uuid.uuid4()).replace("-", ""),
                    "mc": name.strip()
                })

    # 处理被反映人信息
    bfyrs = []
    if pd.notna(row["被反映人姓名"]):
        bfyr_names = str(row["被反映人姓名"]).split(";")
        bfyr_zw = str(row["被反映人职务"]).split(";") if pd.notna(row["被反映人职务"]) else [""] * len(bfyr_names)
        bfyr_zj = str(row["被反映人职级"]).split(";") if pd.notna(row["被反映人职级"]) else [""] * len(bfyr_names)
        bfyr_dwhdz = str(row["被反映人单位或地址"]).split(";") if pd.notna(row["被反映人单位或地址"]) else [""] * len(bfyr_names)
        
        # 确保所有列表长度一致
        max_len = max(len(bfyr_names), len(bfyr_zw), len(bfyr_zj), len(bfyr_dwhdz))
        bfyr_zw.extend([""] * (max_len - len(bfyr_zw)))
        bfyr_zj.extend([""] * (max_len - len(bfyr_zj)))
        bfyr_dwhdz.extend([""] * (max_len - len(bfyr_dwhdz)))
        
        for i, name in enumerate(bfyr_names):
            if name.strip():
                bfyrs.append({
                    "bh": str(uuid.uuid4()).replace("-", ""),
                    "mc": name.strip(),
                    "zw": bfyr_zw[i].strip(),
                    "zj": bfyr_zj[i].strip(),
                    "dwhdz": bfyr_dwhdz[i].strip()
                })

    # 处理初次件信息
    cfj = []
    if pd.notna(row["初次件"]):
        cfj.append({
            "bh": str(uuid.uuid4()).replace("-", ""),
            "wtxx": {
                "yjms": str(row["初次件"])
            }
        })

    return {
        "bh": str(row["信访举报件编号"]),
        "caseIntruction": str(row["原件描述"]),
        "fyrs": fyrs,
        "bfyrs": bfyrs,
        "cfj": cfj
    }

async def send_request(session: aiohttp.ClientSession, data: Dict) -> Dict:
    """发送单个请求"""
    try:
        async with session.post(args.server_url, json=data) as response:
            return await response.json()
    except Exception as e:
        return {"success": "false", "message": str(e), "data": {}}

async def main():
    # 读取Excel文件
    try:
        df = pd.read_excel(args.excel_file, sheet_name="Sheet1")
    except Exception as e:
        print(f"读取Excel文件失败: {str(e)}")
        return

    # 准备测试数据
    test_data_list = []
    for _, row in df.iterrows():
        test_data = prepare_test_data(row)
        test_data_list.append(test_data)

    # 如果测试数据不足，则重复使用现有数据
    while len(test_data_list) < args.num_requests:
        test_data_list.extend(test_data_list[:args.num_requests - len(test_data_list)])

    # 随机选择指定数量的测试数据
    selected_data = random.sample(test_data_list, args.num_requests)

    print(f"开始并发测试，请求数量: {args.num_requests}")
    start_time = time.time()

    # 创建会话并发送并发请求
    async with aiohttp.ClientSession() as session:
        tasks = [send_request(session, data) for data in selected_data]
        results = await asyncio.gather(*tasks)

    end_time = time.time()
    total_time = end_time - start_time

    # 统计结果
    success_count = sum(1 for r in results if r.get("success") == "true")
    fail_count = len(results) - success_count

    print("\n测试结果统计:")
    print(f"总请求数: {len(results)}")
    print(f"成功请求数: {success_count}")
    print(f"失败请求数: {fail_count}")
    print(f"总耗时: {total_time:.2f}秒")
    print(f"平均每个请求耗时: {total_time/len(results):.2f}秒")
    print(f"QPS (每秒查询数): {len(results)/total_time:.2f}")

    # 保存详细结果到文件
    result_file = f"net_concurrent_test_result_{int(time.time())}.json"
    with open(result_file, 'w', encoding='utf-8') as f:
        json.dump({
            "test_config": {
                "num_requests": args.num_requests,
                "server_url": args.server_url,
                "excel_file": args.excel_file
            },
            "summary": {
                "total_requests": len(results),
                "success_count": success_count,
                "fail_count": fail_count,
                "total_time": total_time,
                "avg_time_per_request": total_time/len(results),
                "qps": len(results)/total_time
            },
            "detailed_results": results
        }, f, ensure_ascii=False, indent=2)
    print(f"\n详细结果已保存到: {result_file}")

if __name__ == "__main__":
    asyncio.run(main()) 